#!/usr/bin/python

import sys

import numpy as np

from spectral_clustering_figure import plot_figure
from data import circle_samples, radial_kernel, mutual_knn, rename_clusters

from scipy.sparse import lil_matrix
from scipy.cluster.vq import kmeans2
from scipy.sparse.linalg.eigen.arpack import eigen

def cluster_points(L):
    # sparse eigen is a little bit faster than eig
    evals, evcts = eigen(L, k=15, which="SM")
    #evals, evcts = eig(L)
    evals, evcts = evals.real, evcts.real
    edict = dict(zip(evals, evcts.transpose()))
    evals = sorted(edict.keys())
    # second and third smallest eigenvalue + vector
    Y = np.array([edict[k] for k in evals[1:3]]).transpose()
    res, idx = kmeans2(Y, 3, minit='random')
    return evals[:15], Y, rename_clusters(idx)

def main(args):
    points = circle_samples()
    W = mutual_knn(points)
    G = lil_matrix(W.shape)
    G.setdiag([W.getrow(i).sum() for i in range(W.shape[0])])

    # unnormalized graph Laplacian
    L = G - W
    evals, Y, idx = cluster_points(L)
    
    plot_figure(points, evals, Y, idx)

if __name__ == "__main__":
    main(sys.argv)
